#!/usr/bin/env python
import xml.etree.ElementTree as ET
import urllib

def GetUniprotFile(proteinID):
    for iter in uni_proteinID:  
        uniprot_file = urllib.urlretrieve('http://www.uniprot.org/uniprot/'+iter+'.xml', iter+'.xml')

def GetName(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Get the name of the protein'''
    try:
        name = (protein.find(ns+'protein')
            .find(ns+'recommendedName')
                .findtext(ns+'fullName'))
    except:
        name = (protein.find(ns+'protein')
                .find(ns+'submittedName')
                .findtext(ns+'fullName'))
    return name
    
def GetAccessions(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Get all accession codes for this protein'''
    accession_list = []
    accessions = protein.findall(ns+'accession')
    for accession in accessions:
        accession_list.append(accession.text)
    return accession_list 

def GetGeneName(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Get the gene name'''
    try:
        gene_name = protein.find(ns+'gene').findtext(ns+'name')
    except:
        # Have to invent short names for proteins that don't have one
        try:
            fullname = (protein.find(ns+'protein')
                        .find(ns+'recommendedName')
                        .findtext(ns+'fullName'))
        except:
            fullname = (protein.find(ns+'protein')
                        .find(ns+'submittedName')
                        .findtext(ns+'fullName'))
        gene_name = fullname[0:13]
    return gene_name

def GetSequence(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Retrieve the protein sequence'''
    sequence = protein.findtext(ns+'sequence')
    seq_clean = re.sub("\s+", "", sequence)
    #seq_wrapped = textwrap.fill(seq_clean, 60)
    return seq_clean

def GetOrganism(protein, ns = '{http://uniprot.org/uniprot}'): 
    '''Retrieve the protein organism'''
    organism_name = 'None'
    organism_list = protein.find(ns+'organism').findall(ns+'name')
    for organism in organism_list:  
        if organism.get('type') == 'scientific':
            organism_name = organism.text
    return organism_name

def GetLocation(protein, ns = '{http://uniprot.org/uniprot}'):
    '''return length with start and stop'''
    
    l_begin = (protein.find(ns+'feature').
               find(ns+'location').
               find(ns+'begin').
               get('position'))
    l_end = (protein.find(ns+'feature')
             .find(ns+'location')
             .find(ns+'end')
             .get('position'))
    return int(l_begin), int(l_end)
    

def GetDomain(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Return domains with start and stop positions'''
    domain_dict = {}
    domain_list = []
    domains = protein.findall(ns+'feature')
    # This loop seems cumbersome
    for feature in domains:
        if feature.get('type') == 'domain':
            domain_list.append(feature)
            for domain in domain_list:
                domain_start = (feature.find(ns+'location')
                                .find(ns+'begin')
                                .get('position'))
                domain_stop = (feature.find(ns+'location')
                               .find(ns+'end')
                               .get('position'))
                domain_span = domain_start+'-'+domain_stop
                domain_dict[feature.get('description')]= domain_span
    return domain_dict

def GetLength(protein, ns = '{http://uniprot.org/uniprot}'):
    '''return the sequence length of protein'''
    length = protein.find(ns+'sequence').get('length')
    return int(length)